import numpy as np
import cv2

from ..utils import inference_engine as ie


class SkuRecognizer(object):
    def __init__(self, model_path, mode=ie.MODE_ORT, **kwargs):
        self.res = (112, 336)
        if mode == ie.MODE_ORT:
            self.channel_last = False
            self.engine = ie.ORTEngine(model_path, **kwargs)
        else:
            assert False, 'error: unsupported mode of inference engine.'

    def load_templates(self, feat_path, template_path, thr=0.3):
        self.tpl_ids = np.load(template_path)
        self.tpl_feat_vec = np.load(feat_path)
        self.tpl_feat_vec = self.tpl_feat_vec / np.linalg.norm(self.tpl_feat_vec, axis=1, keepdims=True)
        self.t_tpl_feat_vec = self.tpl_feat_vec.transpose()        
        self.thr = thr
        
    def preprocess(self, imgs):
        t_imgs = [self._transform(img) for img in imgs]
        inp_square = np.stack([x for x in t_imgs], axis=0) 
        data = list()
        data.append(inp_square.astype(np.float32))    
        return data

    def inference(self, inputs):
        feat_vec = self.engine.invoke(inputs)
        return feat_vec[0]

    def extract_feat(self, imgs):
        inputs = self.preprocess(imgs)
        feat_vec = self.inference(inputs)
        return feat_vec

    def classify(self, feat_vec):
        l2 = np.linalg.norm(feat_vec, axis=1, keepdims=True)
        feat_vec = feat_vec / l2
        score_mat = np.dot(feat_vec, self.t_tpl_feat_vec)
        preds = np.argmax(score_mat, axis=1)
        scores = score_mat[np.arange(len(preds)), preds]
        inner_ids = [np.squeeze(self.tpl_ids[i]) for i in preds]
        output_ids = np.array([int(x) for x in inner_ids])       
        return output_ids, scores

    def _transform(self, img):
        img = cv2.resize(img, (self.res[0], self.res[1]))
        img = (img/127.5)-1
        img = img.transpose(2, 0, 1).astype(np.float32)
        return img
